import openai
import os
import requests
import json
import re
# 函数：加载提示词内容
def load_prompt_seed(file_name="Unit/docs/prompt_seed.txt"):
    if os.path.exists(file_name):
        with open(file_name, "r", encoding="utf-8") as file:
            prompt_seed = file.read()
            return prompt_seed
    else:
        raise FileNotFoundError(f"{file_name} 文件未找到。")


def get_completion(prompt):
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer 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',
    }

    data = {
        "model": "deepseek/deepseek-chat-v3-0324:free",
        "prompt": prompt,
    }

    response = requests.post('https://openrouter.ai/api/v1/chat/completions', headers=headers, data=json.dumps(data))
    print("cot的输出：")
    print(response.json())
    # print(response.json()['choices'][0]['text'])
    return response.json()['choices'][0]['text']


def get_code(file_name):
    if os.path.exists(file_name):
        with open(file_name, "r", encoding="utf-8") as file:
            content = file.read()

        # 通过"==================="分割源码和测试用例
        parts = content.split("===================")
        if len(parts) != 2:
            raise ValueError("文件格式不正确，无法分割测试用例和源码。")

        # 上部分是测试用例，下部分是源码
        test_cases = parts[0].strip()
        source_code = parts[1].strip()
        return test_cases, source_code
    else:
        raise FileNotFoundError(f"{file_name} 文件未找到。")

def get_cot(test_function, similar_file):
    # 加载提示词内容
    prompt_seed = load_prompt_seed()
    similar_function_test_cases, similar_function_code = get_code(similar_file)
    prompt_content = (
        f"{prompt_seed}\n\n"
        f"C#待测函数代码：\n{test_function}\n\n"
        "根据这些信息，生成适用于待测函数的测试用例。"
        "生成的代码请用以下格式包裹：```csharp\n ... \n```。保证能让 pattern = r'```csharp\s*(?:\[Fact\])?\s*(.*?)\s*```'识别"
        f"C#相似函数代码：\n{similar_function_code}\n\n"
        f"C#相似函数测试用例：\n{similar_function_test_cases}\n\n"
    )
    get_num = True
    i = 0
    while get_num and i < 3:

        model_output = get_completion(prompt_content)
        # print("输出已保存到 output.txt 文件中。")
        # 正则表达式模式，用于匹配 `csharp` 代码块
        # print(model_output)
        pattern = r'```csharp\s*(?:\[Fact\])?\s*(.*?)\s*```'


        # 使用 re.findall 提取所有匹配的内容，并指定 re.DOTALL 以匹配多行
        csharp_code_blocks = re.findall(pattern, model_output, re.DOTALL)
# 修改后（正确写法）
        print("列表元素数量：" + str(len(csharp_code_blocks)))
        with open("output.txt", "w", encoding="utf-8") as file:
            # 将列表转换为字符串（用双换行分隔每个代码块）
            content = '\n\n'.join(csharp_code_blocks)
            file.write(content)
        print(csharp_code_blocks)
        if csharp_code_blocks != []:
            get_num = False

        i+=1
        prompt_content = (
            "生成的代码请用以下格式包裹：```csharp\n ... \n```。,保证能让 pattern = r'```csharp\s*(?:\[Fact\])?\s*(.*?)\s*```'识别"
         )
    return csharp_code_blocks[0]
